﻿using System.Collections;
using System.Collections.Generic;
using System.IO;
using UnityEngine;
using UnityEngine.Rendering;
using System.Threading;

public class Prefilter : MonoBehaviour {

	const int NumSamples = 1024;

	static float MipLevel2Roughness(int mipLevel, int nMipLevels) {
		var fmipLevel = (float)mipLevel / (nMipLevels - 1);
		var linearRoughness = fmipLevel * fmipLevel;
		return linearRoughness * linearRoughness;
	}

	static float Roughness2MipLevel(float roughness, int nMipLevels) {
		var linearRoughness = Mathf.Sqrt(roughness);
		var fmipLevel = Mathf.Sqrt(linearRoughness);
		fmipLevel = Mathf.Clamp01(fmipLevel);
		return fmipLevel * (nMipLevels - 1);
	}

	[SerializeField]
	ComputeShader m_shader;
	
	[SerializeField]
	Texture MapEnv;

	[SerializeField]
	RenderTexture MapSpecularLD;

	[SerializeField]
	RenderTexture MapDiffuseLD;

	[SerializeField]
	RenderTexture MapDFG;

	// Use this for initialization
	void Start () {
		var kernelID = 0;
		var groupWidth = 0u;
		var groupHeight = 0u;
		var groupDepth = 0u;
		var blockWidth = 0;
		var blockHeight = 0;
		RenderTexture MapLD = null;

		MapSpecularLD.Create();
		kernelID = m_shader.FindKernel("PrefilterSpecularLDMipLevelZero");
		m_shader.SetTexture(kernelID, nameof(MapEnv), MapEnv);
		MapLD = new RenderTexture(MapSpecularLD.width, MapSpecularLD.height, 1, MapSpecularLD.format, RenderTextureReadWrite.Linear);
		MapLD.enableRandomWrite = true;
		MapLD.Create();
		m_shader.SetTexture(kernelID, nameof(MapLD), MapLD);
		for (var FaceIndex = 0; FaceIndex < 6; FaceIndex++) {
			m_shader.SetInt(nameof(FaceIndex), FaceIndex);
			m_shader.GetKernelThreadGroupSizes(kernelID, out groupWidth, out groupHeight, out groupDepth);
			blockWidth = (MapLD.width - 1) / (int)groupWidth + 1;
			blockHeight = (MapLD.height - 1) / (int)groupHeight + 1;
			m_shader.Dispatch(kernelID, blockWidth, blockHeight, 1);
			Graphics.CopyTexture(MapLD, 0, 0, MapSpecularLD, FaceIndex, 0);
		}
		MapLD.Release();

		m_shader.SetInt(nameof(NumSamples), NumSamples);
		kernelID = m_shader.FindKernel("PrefilterSpecularLD");
		m_shader.SetTexture(kernelID, nameof(MapEnv), MapEnv);
		var nMipLevels = (int)Mathf.Floor(Mathf.Log(Mathf.Max(MapSpecularLD.width, MapSpecularLD.height), 2)) + 1;
		for (var mip = 1; mip < nMipLevels; mip++) {
			var Roughness = MipLevel2Roughness(mip, nMipLevels);
			m_shader.SetFloat(nameof(Roughness), Roughness);
			var width = Mathf.Max(MapSpecularLD.width / (1 << mip), 1);
			var height = Mathf.Max(MapSpecularLD.height / (1 << mip), 1);
			MapLD = new RenderTexture(width, height, 1, MapSpecularLD.format, RenderTextureReadWrite.Linear);
			MapLD.enableRandomWrite = true;
			MapLD.Create();
			m_shader.SetTexture(kernelID, nameof(MapLD), MapLD);
			m_shader.GetKernelThreadGroupSizes(kernelID, out groupWidth, out groupHeight, out groupDepth);
			blockWidth = (MapLD.width - 1) / (int)groupWidth + 1;
			blockHeight = (MapLD.height - 1) / (int)groupHeight + 1;
			for (var FaceIndex = 0; FaceIndex < 6; FaceIndex++) {
				m_shader.SetInt(nameof(FaceIndex), FaceIndex);
				m_shader.Dispatch(kernelID, blockWidth, blockHeight, 1);
				Graphics.CopyTexture(MapLD, 0, 0, MapSpecularLD, FaceIndex, mip);
			}
			MapLD.Release();
		}

		MapDiffuseLD.Create();
		kernelID = m_shader.FindKernel("PrefilterDiffuseLD");
		m_shader.SetTexture(kernelID, nameof(MapEnv), MapEnv);
		MapLD = new RenderTexture(MapDiffuseLD.width, MapDiffuseLD.height, 1, RenderTextureFormat.ARGBFloat, RenderTextureReadWrite.Linear);
		MapLD.enableRandomWrite = true;
		MapLD.Create();
		m_shader.SetTexture(kernelID, nameof(MapLD), MapLD);
		m_shader.GetKernelThreadGroupSizes(kernelID, out groupWidth, out groupHeight, out groupDepth);
		blockWidth = (MapLD.width - 1) / (int)groupWidth + 1;
		blockHeight = (MapLD.height - 1) / (int)groupHeight + 1;
		for (var FaceIndex = 0; FaceIndex < 6; FaceIndex ++) {
			m_shader.SetInt(nameof(FaceIndex), FaceIndex);
			m_shader.Dispatch(kernelID, blockWidth, blockHeight, 1);
			Graphics.CopyTexture(MapLD, 0, 0, MapDiffuseLD, FaceIndex, 0);
		}
		MapLD.Release();

		MapDFG.Release();
		MapDFG.enableRandomWrite = true;
		MapDFG.Create();
		kernelID = m_shader.FindKernel("IntegrateDFG");
		m_shader.SetTexture(kernelID, nameof(MapEnv), MapEnv);
		m_shader.SetTexture(kernelID, nameof(MapDFG), MapDFG);
		m_shader.GetKernelThreadGroupSizes(kernelID, out groupWidth, out groupHeight, out groupDepth);
		blockWidth = (MapDFG.width - 1) / (int)groupWidth + 1;
		blockHeight = (MapDFG.height - 1) / (int)groupHeight + 1;
		m_shader.Dispatch(kernelID, blockWidth, blockHeight, 1);
	}
}
